
#include <rtthread.h>
#include <attitude_estimate.h>
#include <MahonyAHRS.h>
#include <math.h>

static rt_sem_t attitude_estimate_sem = RT_NULL; // 姿态估计的信号量，用于同步传感器读取线程和姿态解算线程
static rt_sem_t attitude_read_sem = RT_NULL;     // 姿态角读取的信号量，用于同步姿态解算线程更新姿态角和其它线程读取姿态角

static struct sensor_estimate_data *sensor_estimate_buf[SENSOR_STORE_MAX]; // 没有经过滤波的存储：加速度、角速度、地磁数据
static int store_pos = 0;                                                  // 存储/更新的当前序号：0~SENSOR_STORE_MAX-1

static struct sensor_estimate_data filter_data; // 经过滤波后的一组：加速度、角速度、地磁数据
static enum attitude_estimate_kind IMU_or_AHRS; // 姿态融合是六轴的IMU(加速度+陀螺仪),还是九轴的AHRS(加速度+陀螺仪+地磁)

static struct attitude_euler current_estimate_euler;                 /* 当前估计的欧拉角 */
static struct attitude_euler_angular current_estimate_euler_angular; /* 当前估计的欧拉角速度 */

int get_attitude(struct attitude_euler *attitude, struct attitude_euler_angular *attitude_angular)
{
    if (!attitude || !attitude_angular)
        return RT_EINVAL;

    rt_sem_take(attitude_read_sem, RT_WAITING_FOREVER);
    *attitude = current_estimate_euler;
    *attitude_angular = current_estimate_euler_angular;

    return 0;
}

/* 根据连续性原理，修正atan以及asin求出来的姿态角 */
static int modify_estimate_attitude(struct attitude_euler *last_attitude)
{
    /*因为atan值域在[-pi/2 ~ pi/2]，需要将其扩展，根据上一时刻与当前时刻的角度连续性判断 */
    /* 如果需要的话修正偏航角 */
    if (last_attitude->yaw > -90.f / RAD2DEG && last_attitude->yaw < -60.f / RAD2DEG &&
        current_estimate_euler.yaw > 0)
        current_estimate_euler.yaw -= 180.f / RAD2DEG;
    if (last_attitude->yaw > 60.f / RAD2DEG && last_attitude->yaw < 90.f / RAD2DEG &&
        current_estimate_euler.yaw < 0)
        current_estimate_euler.yaw += 180.f / RAD2DEG;
    return 0;
}

/// @brief 接收经过滤波的5组：加速度、角速度、地磁数据，估计输出/更新姿态角
/// @param attitude
/// @param filter_data
static void estimate_attitude_update()
{
    float acce_x, acce_y, acce_z;
    float gyro_x, gyro_y, gyro_z;
    float mag_x, mag_y, mag_z;

    acce_x = filter_data.acce->data.acce.x;
    acce_y = filter_data.acce->data.acce.y;
    acce_z = filter_data.acce->data.acce.z;

    gyro_x = filter_data.gyro->data.gyro.x;
    gyro_y = filter_data.gyro->data.gyro.y;
    gyro_z = filter_data.gyro->data.gyro.z;

    mag_x = filter_data.mag->data.mag.x;
    mag_y = filter_data.mag->data.mag.y;
    mag_z = filter_data.mag->data.mag.z;

#if ATTITUDE_ESTIMATE_USING_MAHONY
    if (IMU_or_AHRS == IMU)
        MahonyAHRSupdateIMU(gyro_x, gyro_y, gyro_z, acce_x, acce_y, acce_z);
    else if (IMU_or_AHRS == AHRS)
        MahonyAHRSupdate(gyro_x, gyro_y, gyro_z, acce_x, acce_y, acce_z, mag_x, mag_y, mag_z);

    struct attitude_euler last_attitude = current_estimate_euler;
    MahonyAHRS_ResultFromQuant2Eular(&(current_estimate_euler.roll), &(current_estimate_euler.pitch), &(current_estimate_euler.yaw));

    modify_estimate_attitude(&last_attitude);

    rt_sem_release(attitude_read_sem);
#elif ATTITUDE_ESTIMATE_USING_EKF

#elif ATTITUDE_ESTIMATE_USING_SEPERATE
    current_estimate_euler.pitch = atanf(1.0f * filter_data.acce->x / filter_data.acce->z);
    current_estimate_euler.roll = atanf(1.0f * filter_data.acce->y / filter_data.acce->z);
    current_estimate_euler.yaw = atan2f((filter_data.mag->x * cosf(attitude->pitch) +
                                         filter_data.mag->y * sinf(attitude->roll) * sinf(attitude->pitch) +
                                         filter_data.mag->z * cosf(attitude->roll) * sinf(attitude->pitch)),
                                        (filter_data.mag->y * cosf(attitude->roll) - filter_data.mag->z * sinf(attitude->roll)));
    rt_sem_release(attitude_read_sem);
#endif
    /* 根据角速度和欧拉角速度关系，求出欧拉角速度 */

    /* 方式一：近似表达 */
    // current_estimate_euler_angular.roll_angular = gyro_x;
    // current_estimate_euler_angular.pitch_angular = gyro_y;
    // current_estimate_euler_angular.yaw_angular = gyro_z;

    /* 方式二：计算 */
    current_estimate_euler_angular.pitch_angular = cosf(current_estimate_euler.roll) * gyro_y - sinf(current_estimate_euler.roll) * gyro_z;
    current_estimate_euler_angular.roll_angular = gyro_x + sinf(current_estimate_euler.roll) * tanf(current_estimate_euler.pitch) * gyro_y + cosf(current_estimate_euler.roll) * tanf(current_estimate_euler.pitch) * gyro_z;
    current_estimate_euler_angular.yaw_angular = sinf(current_estimate_euler.roll) / cosf(current_estimate_euler.pitch) * gyro_y + cosf(current_estimate_euler.roll) / cosf(current_estimate_euler.pitch) * gyro_z;
}

/// @brief 对目前的5组传感器数据进行均值，随后使用算法获得姿态角
/// @param filter_data
static void filter_sensor_data()
{
    filter_data.acce->data.acce.x = 0;
    filter_data.acce->data.acce.y = 0;
    filter_data.acce->data.acce.z = 0;

    filter_data.gyro->data.gyro.x = 0;
    filter_data.gyro->data.gyro.y = 0;
    filter_data.gyro->data.gyro.z = 0;

    filter_data.mag->data.mag.x = 0;
    filter_data.mag->data.mag.y = 0;
    filter_data.mag->data.mag.z = 0;

    for (int i = 0; i < SENSOR_STORE_MAX; i++)
    {
        filter_data.acce->data.acce.x += sensor_estimate_buf[i]->acce->data.acce.x;
        filter_data.acce->data.acce.y += sensor_estimate_buf[i]->acce->data.acce.y;
        filter_data.acce->data.acce.z += sensor_estimate_buf[i]->acce->data.acce.z;

        filter_data.gyro->data.gyro.x += sensor_estimate_buf[i]->gyro->data.gyro.x;
        filter_data.gyro->data.gyro.y += sensor_estimate_buf[i]->gyro->data.gyro.y;
        filter_data.gyro->data.gyro.z += sensor_estimate_buf[i]->gyro->data.gyro.z;

        filter_data.mag->data.mag.x += sensor_estimate_buf[i]->mag->data.mag.x;
        filter_data.mag->data.mag.y += sensor_estimate_buf[i]->mag->data.mag.y;
        filter_data.mag->data.mag.z += sensor_estimate_buf[i]->mag->data.mag.z;
    }

    filter_data.acce->data.acce.x /= SENSOR_STORE_MAX;
    filter_data.acce->data.acce.y /= SENSOR_STORE_MAX;
    filter_data.acce->data.acce.z /= SENSOR_STORE_MAX;

    filter_data.gyro->data.gyro.x /= SENSOR_STORE_MAX;
    filter_data.gyro->data.gyro.y /= SENSOR_STORE_MAX;
    filter_data.gyro->data.gyro.z /= SENSOR_STORE_MAX;

    filter_data.mag->data.mag.x /= SENSOR_STORE_MAX;
    filter_data.mag->data.mag.y /= SENSOR_STORE_MAX;
    filter_data.mag->data.mag.z /= SENSOR_STORE_MAX;
}

/// @brief 姿态估计工作线程
/// @param params
static void attitude_estimate_thread(void *params)
{
    while (1)
    {
        /* 当更新加速度/陀螺仪/地磁数据之后，另外一个线程调用update_attitude_sensor释放信号量attitude_estimate_sem
            此时姿态估计线程就调用rt_sem_take获取信号量
         */
        rt_sem_take(attitude_estimate_sem, RT_WAITING_FOREVER);

        store_pos++;
        if (store_pos == SENSOR_STORE_MAX)
        {
            filter_sensor_data();
            estimate_attitude_update();
            store_pos = 0;
        }
    }
}

int start_attitude_estimate()
{
    /* 创建信号量 */
    attitude_estimate_sem = rt_sem_create("attitude_sem", 0, RT_IPC_FLAG_PRIO);
    if (attitude_estimate_sem == RT_NULL)
    {
        rt_kprintf("create attitude estimate semophore failed\n");
        return RT_ERROR;
    }

    attitude_read_sem = rt_sem_create("attitude_rw_sem", 0, RT_IPC_FLAG_PRIO);
    if (attitude_read_sem == RT_NULL)
    {
        rt_kprintf("create attitude read write semophore failed\n");
        return RT_ERROR;
    }

    /* 申请用于存储传感器多组数据的buf */
    for (int i = 0; i < SENSOR_STORE_MAX; i++)
    {
        sensor_estimate_buf[i] = rt_malloc(sizeof(struct sensor_estimate_data));
        if (sensor_estimate_buf[i] == RT_NULL)
        {
            rt_kprintf("attitude estimate alloc memory for estimate sensor data failed\n");
            return -3;
        }
        sensor_estimate_buf[i]->acce = rt_malloc(sizeof(sensor_data_t));
        sensor_estimate_buf[i]->gyro = rt_malloc(sizeof(sensor_data_t));
        sensor_estimate_buf[i]->mag = rt_malloc(sizeof(sensor_data_t));

        if (sensor_estimate_buf[i]->acce == RT_NULL || sensor_estimate_buf[i]->gyro == RT_NULL || sensor_estimate_buf[i]->mag == RT_NULL)
        {
            rt_kprintf("attitude estimate alloc memory for sensor data failed\n");
            return RT_ENOMEM;
        }

        rt_memset(sensor_estimate_buf[i]->acce, 0, sizeof(sensor_data_t));
        rt_memset(sensor_estimate_buf[i]->gyro, 0, sizeof(sensor_data_t));
        rt_memset(sensor_estimate_buf[i]->mag, 0, sizeof(sensor_data_t));
    }

    /* 申请用于存储传感器滤波数据 */
    filter_data.acce = rt_malloc(sizeof(sensor_data_t));
    filter_data.gyro = rt_malloc(sizeof(sensor_data_t));
    filter_data.mag = rt_malloc(sizeof(sensor_data_t));
    if (filter_data.acce == RT_NULL || filter_data.gyro == RT_NULL || filter_data.mag == RT_NULL)
    {
        rt_kprintf("attitude estimate alloc memory for sensor filter data failed\n");
        return RT_ENOMEM;
    }

    rt_memset(filter_data.acce, 0, sizeof(sensor_data_t));
    rt_memset(filter_data.gyro, 0, sizeof(sensor_data_t));
    rt_memset(filter_data.mag, 0, sizeof(sensor_data_t));

    /* 创建姿态估计线程 */
    rt_thread_t attitude_est = rt_thread_create(
        "attitude_estimate",
        attitude_estimate_thread,
        RT_NULL,
        ATTITUDE_ESTIMATE_THREAD_STACK_SIZE,
        ATTITUDE_ESTIMATE_THREAD_PRIORITY,
        ATTITUDE_ESTIMATE_THREAD_TICK);

    if (attitude_est)
    {
        rt_thread_startup(attitude_est);
        return 0;
    }

    rt_kprintf("start attitude estimate failed\n");
    return -1;
}

int update_attitude_sensor(sensor_data_t *acce, sensor_data_t *gyro, sensor_data_t *mag)
{
    if (acce == RT_NULL || gyro == RT_NULL)
        return RT_EINVAL;

    IMU_or_AHRS = (mag == RT_NULL) ? IMU : AHRS;

    if (sensor_estimate_buf[store_pos]->acce)
        rt_memcpy(sensor_estimate_buf[store_pos]->acce, acce, sizeof(sensor_data_t));

    if (sensor_estimate_buf[store_pos]->gyro && gyro)
        rt_memcpy(sensor_estimate_buf[store_pos]->gyro, gyro, sizeof(sensor_data_t));

    if (sensor_estimate_buf[store_pos]->mag && mag)
        rt_memcpy(sensor_estimate_buf[store_pos]->mag, mag, sizeof(sensor_data_t));

    rt_sem_release(attitude_estimate_sem);

    return 0;
}
